import torch
from models.backbone import ModelWrapper, INPUT_SIZE
from models.classifier import Classifier
from models.loss import MultiCBFocalLoss, CB_Weight
from trainer.trainer import ClasificationTrainer
from utils.consts import TASK_CONF_PATH


class Atrainer(ClasificationTrainer):
    def __init__(self, task_id: str, conf: dict, title_icon: str) -> None:
        super().__init__(task_id, conf, title_icon)

    def prepare(self):
        from dataproc.adjust import get_loaders
        self.tr_loader, self.vl_loader \
            = get_loaders(
                name=self.dataset_name,
                distr=self.distr,
                batch_size=self.batch,
                input_size=INPUT_SIZE)

        self.distr = self.tr_loader.dataset.distr
        self.cls_dict = self.tr_loader.dataset.cls_dict
        self.num_cls = len(self.distr)

        self.net = ModelWrapper(name=self.backbone, num_cls=self.num_cls)
        # self.lossfunc = MultiCBFocalLoss(
        #     self.num_cls, gamma=0.5, alpha=CB_Weight(self.distr).to(self.device))
        self.lossfunc = torch.nn.CrossEntropyLoss()
        super().prepare()


class Btrainer(ClasificationTrainer):
    def __init__(self, task_id: str, conf: dict, title_icon: str) -> None:
        super().__init__(task_id, conf, title_icon)
        self.model_path = f"{TASK_CONF_PATH}/{conf['model']}"
        self.methods = conf['method']
        self.gan_path = None
        if 'GAN' in self.methods:
            self.gan_path = f"{TASK_CONF_PATH}/{conf['gan']}"

    def prepare(self):
        from dataproc.augment import get_loaders
        self.tr_loader, self.vl_loader = get_loaders(
            name=self.dataset_name,
            distr=self.distr,
            batch_size=self.batch,
            backbone=self.backbone,
            model_pth=self.model_path,
            methods=self.methods,
            gan_path=self.gan_path,
            device=self.device)

        self.distr = [self.tr_loader.dataset.org_distr,
                      self.tr_loader.dataset.distr]
        self.cls_dict = self.tr_loader.dataset.cls_dict
        self.feature_num = self.tr_loader.dataset.feature_num
        self.num_cls = len(self.cls_dict)

        self.net = Classifier(in_feats=self.feature_num, num_cls=self.num_cls)
        self.lossfunc = MultiCBFocalLoss(
            self.num_cls, gamma=0.5, alpha=CB_Weight(self.distr).to(self.device))
        super().prepare()
